#pragma once

#include <iostream>
using namespace std;

enum Colour
{
    BLACK,
    RED
};

template <class K, class V>
struct RBTreeNode
{
    RBTreeNode<K, V>* _left;
    RBTreeNode<K, V>* _right;
    RBTreeNode<K, V>* _parent;
    Colour _col;
    pair<K, V> _kv;

    RBTreeNode(const pair<K, V>& kv)
        : _left(nullptr)
        , _right(nullptr)
        , _parent(nullptr)
        , _col(RED)
        , _kv(kv)
    {}
};

template <class K, class V>
class RBTree
{
    typedef RBTreeNode<K, V> Node;
private:
    Node* _root;
public:
    RBTree(): _root(nullptr)
    {}
    bool Insert(const pair<K, V>& kv)
    {
        if(_root == nullptr)
        {
            _root = new Node(kv);
            _root->_col = BLACK;
            return true;
        }

        Node* parent = nullptr;
        Node* cur = _root;

        while(cur != nullptr)
        {
            if(cur->_kv.first < kv.first)
            {
                parent = cur;
                cur = cur->_right;
            }
            else if(cur->_kv.first > kv.first)
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return false;
            }
        }
        cur = new Node(kv);
        // cur->_col = RED;
        if(parent->_kv.first > kv.first)
        {
            parent->_left = cur;
            cur->_parent = parent;
        }
        else
        {
            parent->_right = cur;
            cur->_parent = parent;
        }

        //balance control
        while(parent != nullptr && parent->_col == RED)
        {
            Node* grandfather = parent->_parent;
            if(parent == grandfather->_left)
            {
                Node* uncle = grandfather->_right;
                //情况一：叔叔存在且为红，进行变色处理，并继续向上处理
                if(uncle && uncle->_col == RED)
                {
                    parent->_col = uncle->_col = BLACK;
                    grandfather->_col = RED;

                    cur = grandfather;
                    parent = cur->_parent;
                }
                else    //uncle不存在 or 存在且为黑
                {
                    if(cur == parent->_left)    //情况二：单旋 + 变色
                    {
                        RotateR(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    else    //情况三：双旋 + 变色
                    {
                        RotateL(parent);
                        RotateR(grandfather);
                        grandfather->_col = RED;
                        cur->_col = BLACK;
                    }
                    break;
                }
            }
            else
            {
                Node* uncle = grandfather->_left;
                if(uncle && uncle->_col == RED)
                {
                    parent->_col = BLACK;
                    uncle->_col = BLACK;
                    grandfather->_col = RED;

                    cur = grandfather;
                    parent = cur->_parent;
                }
                else
                {
                    if(cur == parent->_right)
                    {
                        RotateL(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    else
                    {
                        RotateR(parent);
                        RotateL(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    break;
                }
            }
        }
        _root->_col = BLACK;
        return true;
    }

    void RotateL(Node* parent)
    {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        parent->_right = subRL;
        if(subRL != nullptr)
        {
            subRL->_parent = parent;
        }
        Node* ppNode = parent->_parent;
        subR->_left = parent;
        parent->_parent = subR;

        if(parent == _root)
        {
            _root = subR;
            subR->_parent = nullptr;
        }
        else
        {
            if(ppNode->_left == parent)
            {
                ppNode->_left = subR;
                subR->_parent = ppNode;
            }
            else
            {
                ppNode->_right = subR;
                subR->_parent = ppNode;
            }
        }
        
    }

    void RotateR(Node* parent)
    {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        parent->_left = subLR;
        if(subLR != nullptr)
        {
            subLR->_parent = parent;
        }
        Node* ppNode = parent->_parent;
        subL->_right = parent;
        parent->_parent = subL;

        if(parent == _root)
        {
            _root = subL;
            subL->_parent = nullptr;
        }
        else
        {
            if(ppNode->_left == parent)
            {
                ppNode->_left = subL;
            }
            else
            {
                ppNode->_right = subL;
            }
            subL->_parent = ppNode;
        }
    }

    bool IsBalance()
    {
        if(_root == nullptr)
        {
            return true;
        }
        if(_root->_col == RED)
        {
            cerr << "error: root is RED!" << endl;
            return false;
        }

        //算出最左路径的黑色节点作为一个基准值
        Node* cur = _root;
        int benchmark = 0;
        while(cur)
        {
            if(cur->_col == BLACK)
            {
                benchmark++;
            }
            cur = cur->_left;
        }

        return CheckRED_RED(_root)
            && CheckBlackNum(_root, 0, benchmark);
    }

    void InOrder()
    {
        _InOrder(_root);
    }
  private:
    void _InOrder(Node* root)
    {
        if(root == nullptr)
        {
            return;
        }
        _InOrder(root->_left);
        cout << root->_kv.first << " : " << root->_kv.second << " " << endl;
        _InOrder(root->_right);
    }
    bool CheckRED_RED(Node* cur)
    {
        if(cur == nullptr)
        {
            return true;
        }

        if(cur->_col == RED && cur->_parent->_col == RED)
        {
            cerr << "error: 连续红节点！" << endl;
            return false; 
        }

        return CheckRED_RED(cur->_left)
            && CheckRED_RED(cur->_right);
    }
    bool CheckBlackNum(Node* cur, int blackNum, int benchmark)
    {
        if(cur == nullptr)
        {
            if(blackNum != benchmark)
            {
                cerr << "error: 每条路径上的黑色节点不完全相同！" << endl;
                return false;
            }
            return true;
        }
        if(cur->_col == BLACK)
        {
            blackNum++;
        }
        return CheckBlackNum(cur->_left, blackNum, benchmark)
            && CheckBlackNum(cur->_right, blackNum, benchmark);
    }
};
